import re

import pandas as pd
from tqdm import tqdm

from iflytech_assistant.ugc.role import Example, RoleMeta
from iflytech_assistant.ugc.role_play import generate

df = pd.read_excel(".vscode/ugc_online/初音未来.xlsx")


def run_single(name):
    if name not in df["助手名称"].values:
        raise ValueError(f"未找到角色名: {name}")
    profile = df[df["助手名称"] == name]["助手简介"].values[0]
    introduction = df[df["助手名称"] == name]["助手描述"].values[0]
    category = df[df["助手名称"] == name]["用户创建时分类"].values[0]

    examples = df[df["助手名称"] == name]["输入输出示例"].values[0]

    if isinstance(examples, str):
        examples = examples.strip()
        examples = re.sub(r"\n\r+", "\n", examples)
        examples = examples.split("\n")
        examples = [
            item for item in examples if re.sub(r"输[入出]：", "", item).strip()
        ]
    else:
        examples = []
    lines = df[df["助手名称"] == name]["用户输入的语录"].values[0]
    if isinstance(lines, str):
        lines = lines.strip().split("\n")
    else:
        lines = []

    role_meta = RoleMeta(
        name=name,
        profile=profile,
        introduction=introduction,
        category=category,
        lines=lines,
    )
    # 新Prompt输出
    output = ""
    for example in examples:
        if re.search(r"输入：", example) is None:
            continue
        user_input = example.replace("输入：", "")
        example: Example = generate(
            role_meta, user_input=user_input, _print=False, multi_results=False
        )
        output += f"用户输入：{example.user_input}\n"
        output += f"助手输出：{example.role_output}\n"

    return output


if __name__ == "__main__":

    outputs = []
    for name in tqdm(df["助手名称"].values):
        try:
            output = run_single(name)
        except Exception as e:
            output = ""

        outputs.append(output)

    df["新Prompt输出"] = outputs
    df.to_excel(".vscode/ugc_online/初音未来2.xlsx", index=False)
